
import { Chroma } from "@langchain/community/vectorstores/chroma";
import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { PDFLoader } from "langchain/document_loaders/fs/pdf";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";

const collectionName = "hcy-traingo"
const embeddings = new OllamaEmbeddings({
    model: "milkey/m3e", // default value
    baseUrl: "http://117.72.38.226:11434", // default value
    requestOptions: {
    useMMap: true,
    numThread: 6,
    // numGpu: 1,
    },
});
const loader = new PDFLoader("./doc/traingo学习管理系统.pdf");
const docs = await loader.load();
const textSplitter = new RecursiveCharacterTextSplitter(
    {
      chunkSize: 500,
      chunkOverlap: 300,
    }
    );
const splitDocs = await textSplitter.splitDocuments(docs);
await Chroma.fromDocuments(
  splitDocs,
  embeddings,
  {
      collectionName: collectionName, // 
      url: "http://117.72.38.226:8000", // Optional, will default to this value
  }
);
const vectorStore = new Chroma(
    embeddings,
    {
      collectionName: collectionName,
      url: "http://117.72.38.226:8000",
    }
)
const result = await vectorStore.similaritySearch("金币怎么获取", 3);
console.log(result)
